import numpy as np
import bz2

# 转为字节：每8位一组
def bin_str_to_bytes(s: str):
    s = s.ljust((len(s) + 7) // 8 * 8, "0")  # 补齐到8的倍数
    return bytes([int(s[i : i + 8], 2) for i in range(0, len(s), 8)])

# 转为二进制字符串：将字节流还原为 '0' 和 '1' 组成的字符串
def bytes_to_bin_str(data: bytes) -> str:
    return ''.join(f'{byte:08b}' for byte in data)

def compressMethod(lsbs: list[np.uint8]) -> bytes:
    # 使用 bz2 压缩
    lsbsStr = "".join(lsbs)  # 拼接成字符串
    # 转为字节
    data = bin_str_to_bytes(lsbsStr)
    compressed = bz2.compress(data)
    # print(f"bz2 压缩前: {len(data)} 字节")
    # print(f"bz2 压缩后: {len(compressed)} 字节")
    # print(f"压缩比: {len(compressed) / len(data):.2f}")
    return compressed

# 解压：将 bz2 压缩数据还原为原始的 LSB 位列表（list[str]）
def decompressMethod(compressed: bytes) -> list[str]:
    # 使用 bz2 解压
    decompressed_data = bz2.decompress(compressed)
    # 将 bytes 转回二进制字符串
    bin_str = bytes_to_bin_str(decompressed_data)
    # 转为 list['0', '1', ...]
    return list(bin_str)

    
if __name__ == "__main__":
    # 示例：0 占 80%，1 占 20%
    lsbs = ["0"] * 800 + ["1"] * 200
    print(f"原始数据长度: {len(lsbs)} bits")

    # 压缩
    compressed = compressMethod(lsbs)
    print(f"压缩后数据长度: {len(compressed)} bytes")

    # 解压
    decompressed = decompressMethod(compressed)
    print(f"解压后数据长度: {len(decompressed)} bits")

    # 验证一致性
    if lsbs == decompressed:
        print("✅ 压缩 -> 解压 成功，数据一致！")
    else:
        print("❌ 数据不一致！")

    # 压缩率（按字节计算）
    original_bytes_length = (len(lsbs) + 7) // 8  # 原始二进制串占用的字节数
    compression_ratio = len(compressed) / original_bytes_length if original_bytes_length > 0 else 0
    print(f"压缩率: {compression_ratio:.2f} ({len(compressed)}/{original_bytes_length})")